package com.achieve.ai.config;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeEmbeddingStore;
import dev.langchain4j.store.embedding.pinecone.PineconeServerlessIndexConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
@Slf4j
public class EmbeddingStoreConfig {
    @Autowired
    private EmbeddingModel embeddingModel;

    @Value("${milvus.host}")
    private String host;
    @Value("${milvus.port}")
    private Integer port;

    /**
     * PINECONE
     *
     * @return
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        //创建向量存储
        EmbeddingStore<TextSegment> embeddingStore = PineconeEmbeddingStore.builder()
                .apiKey(System.getenv("PINECONE_API_KEY"))
                .index("xiaozhi-index")//如果指定的索引不存在，将创建一个新的索引
                .nameSpace("xiaozhi-namespace") //如果指定的名称空间不存在，将创建一个新的名称空间
                .createIndex(PineconeServerlessIndexConfig.builder()
                        .cloud("AWS") //指定索引部署在 AWS 云服务上。
                        .region("us-east-1") //指定索引所在的 AWS 区域为 us-east-1。
                        .dimension(embeddingModel.dimension()) //指定索引的向量维度，该维度与 embeddedModel 生成的向量维度相同。
                        .build())
                .build();
        return embeddingStore;
    }

//    @Bean
//    public EmbeddingStore createEmbeddingStore() {
//        log.info("==========开始创建Milvus的Collection");
//        MilvusEmbeddingStore store = MilvusEmbeddingStore.builder()
//                .host(host)
//                .port(port)
//                .collectionName("langchain_01")
//                .dimension(1024)
//                .indexType(IndexType.FLAT)
//                .metricType(MetricType.COSINE)
////                .username("username")
////                .password("password")
//                .consistencyLevel(ConsistencyLevelEnum.EVENTUALLY)
//                .autoFlushOnInsert(true)
//                .idFieldName("id")
//                .textFieldName("text")
//                .metadataFieldName("metadata")
//                .vectorFieldName("vector")
//                .build();
//        log.info("==========创建Milvus的Collection完成");
//        return store;
//    }

}
